{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第一步：测试模型是否可用"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "model_name = \"/e/Resources/LLM/huggingface/DeepSeek-R1-Distill-Qwen-1.5B\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
    "# model = AutoModelForCausalLM.from_pretrained(model_name).to(\"cuda\")\n",
    "print(\"-----模型加载成功-----\")\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第二步：制作数据集\n"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from data_prepare import samples\n",
    "import json\n",
    "\n",
    "with open(\"datasets.jsonl\", \"w\", encoding=\"utf-8\") as f:\n",
    "    for s in samples:\n",
    "        json_line = json.dumps(s, ensure_ascii=False)\n",
    "        f.write(json_line + \"\\n\")\n",
    "    else:\n",
    "        print(\"prepare data finished\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第三步： 准备训练集和测试集"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"json\", data_files={\"train\": \"datasets.jsonl\"}, split=\"train\")\n",
    "print(\"数据数量:\", len(dataset))\n",
    "\n",
    "train_test_split = dataset.train_test_split(test_size=0.1)  # 45:5\n",
    "train_dataset = train_test_split[\"train\"]\n",
    "eval_dataset = train_test_split[\"test\"]\n",
    "print(f\"train dataset len: {len(train_dataset)}\")\n",
    "print(f\"test dataset len: {len(eval_dataset)}\")\n",
    "\n",
    "print(\"-----完成训练数据的准备工作---\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第四步： 编写tokenizer处理工具"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def tokenizer_function(many_samples):\n",
    "    texts = [f\"{prompt}\\n{completion}\" for prompt, completion in\n",
    "             zip(many_samples[\"prompt\"], many_samples[\"completion\"])]\n",
    "    tokens = tokenizer(texts, truncation=True, max_length=512, padding=\"max_length\")\n",
    "    tokens[\"labels\"] = tokens[\"input_ids\"].copy()\n",
    "    return tokens\n",
    "\n",
    "\n",
    "tokenized_train_dataset = train_dataset.map(tokenizer_function, batched=True)\n",
    "tokenized_eval_dataset = eval_dataset.map(tokenizer_function, batched=True)\n",
    "\n",
    "print(\"----完成tokenizing---\")\n",
    "print(\"----train data---\")\n",
    "print(tokenized_train_dataset[0])\n",
    "print(\"----eval data---\")\n",
    "print(tokenized_eval_dataset[0])"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第五步：量化设置（失败，macbook不支持bitsandbytes）"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from transformers import BitsAndBytesConfig #需要gpu支持cuda\n",
    "\n",
    "quantization_config = BitsAndBytesConfig(load_in_8bit=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quantization_config, device_map=\"auto\")\n",
    "print(\"-----已经完成量化模型的加载-----\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第六步：lora微调设置"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "\n",
    "lora_config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, task_type=TaskType.CAUSAL_LM)\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()\n",
    "print(\"----lora微调设置完毕-----\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第七步：设置训练参数"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./sft_model\",\n",
    "    num_train_epochs=10,\n",
    "    per_device_train_batch_size=4,\n",
    "    gradient_accumulation_steps=8,\n",
    "    fp16=True,\n",
    "    logging_steps=10,\n",
    "    save_steps=100,\n",
    "    eval_strategy=\"steps\",\n",
    "    eval_steps=10,\n",
    "    learning_rate=3e-5,\n",
    "    logging_dir=\"./logs\",\n",
    "    run_name=\"deepseek-r1-distill-finetune\"\n",
    ")\n",
    "print(\"----设置训练参数完毕-----\")\n",
    "# 定义训练器\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_train_dataset,\n",
    "    eval_dataset=tokenized_eval_dataset\n",
    ")\n",
    "print(\"---开始训练-----\")\n",
    "trainer.train()\n",
    "print(\"--训练完成--\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 第八步： 保存lora模型"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "#保存lora模型\n",
    "save_path =\"./saved models\"\n",
    "model.save_pretrained(save_path)\n",
    "tokenizer.save_pretrained(save_path)\n",
    "print(\"---lora模型已经保存---\")\n",
    "\n",
    "# 保存全量模型\n",
    "final_save_path =\"./final_saved_path\"\n",
    "from peft import PeftModel\n",
    "base_model = AutoModelForCausalLM.from_pretrained(model_name)\n",
    "model = PeftModel.from_pretrained(base_model,save_path)\n",
    "model = model.merge_and_unload()\n",
    "model.save_pretrained(final_save_path)\n",
    "tokenizer.save_pretrained(final_save_path)"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rag",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
